from graph import *
import numpy as np
import random
import copy
from bezier import transition as bt


class agent:
    def __init__(self, pos=vector2d()):
        self.position = pos
        self.initposition=None
        self.initvertex=None
        self.id = self.register()
        self.lastdirection = vector2d()
        self.direction = vector2d()
        self.vertexes = {}
        self.treeneighbors={}

    def setposition(self,pos):
        lastpos=self.position
        self.position = pos
        key=str(pos.x*1000+pos.y)
        self.vertexes[key]=stcvertex(x=pos.x,y=pos.y)
        print key
        print len(self.vertexes)
        if(instance.graph.vertexes[key].state == "reached"):
            instance.graph.vertexes[key].overreached=True
        instance.graph.vertexes[key].state = "reached"
        nb = instance.graph.getneighbors(self.position)
        ev=None

        treeflag=True
        for n in nb:
            k = str(n.center.x*1000 + n.center.y)
            if not self.vertexes.has_key(k):
                self.treeneighbors[k]=instance.graph.vertexes[k]
            else:
                if treeflag:
                    self.vertexes[str(pos.x * 1000 + pos.y)].edges[str(n.center.x * 1000 + n.center.y)] = self.vertexes[
                        str(n.center.x * 1000 + n.center.y)]
                    self.vertexes[str(n.center.x * 1000 + n.center.y)].edges[str(pos.x * 1000 + pos.y)] = self.vertexes[
                        str(pos.x * 1000 + pos.y)]
                    ev = self.vertexes[str(n.center.x * 1000 + n.center.y)].center
                    treeflag=False
        if self.treeneighbors.has_key(key):
            self.treeneighbors.pop(key)
        if self.initposition is None:
            self.initposition=vector2d(pos.x-0.25,pos.y-0.25)
            self.initvertex=self.vertexes[key]
        return ev

    def unreachedneighbor(self):
        rn = {}
        for n in self.treeneighbors:
            if not self.treeneighbors[n].state == "reached":
                k = str(self.treeneighbors[n].center.x*1000+self.treeneighbors[n].center.y)
                rn[k]=self.treeneighbors[k]
        return rn


    def sortvertex(self, res=1,index=0):
        keys=self.vertexes.keys()
        hc={}
        vc={}
        for k in keys:
            p=self.vertexes[k].center
            if not hc.has_key(str(p.x)):
                hc[str(p.x)] = 1
            else:
                hc[str(p.x)] = hc[str(p.x)] + 1
            if not vc.has_key(str(p.y)):
                vc[str(p.y)] = 1
            else:
                vc[str(p.y)] = vc[str(p.y)]+1
        # primary movement direction is vertical
        branches={}
        direction=1
        L=sorted(hc.items(),key=lambda x:x[1],reverse=True)
        hcm=0
        vcm=0
        for l in L[:2]:
            hcm=hcm+l[1]
        L = sorted(vc.items(), key=lambda x: x[1], reverse=True)
        for l in L[:2]:
            vcm=vcm+l[1]
        # primary movement direction is horizontal
        if(vcm>hcm):
            direction=0
            for k in vc.keys():
                branches[str(k)]=[]
        else:
            for k in hc.keys():
                branches[str(k)]=[]

        for k in keys:
            p = self.vertexes[k]
            p.edges={}
            if direction>0:
                branches[str(p.center.x)].append(p)
            else:
                branches[str(p.center.y)].append(p)
        for k in branches.keys():
            b = branches[k]
            #print b
            if direction>0:
                branches[k] = sorted(b,key=lambda t : t.center.y)
            else:
                branches[k] = sorted(b, key=lambda t: t.center.x)
        #branches=sorted(branches)
        for k in branches:
            b=branches[k]
            adjb=False
            pflag=False
            print b
            for bi in b:
                if direction>0:
                    if adjb==False:
                        firstcon = False
                        bk = str((bi.center.x - res) * 1000 + bi.center.y)
                        if self.vertexes.has_key(bk):
                            adjb = True
                            firstcon = True
                            bi.edges[bk]=self.vertexes[bk]
                            self.vertexes[bk].edges[str(bi.center.x * 1000 + bi.center.y)] = bi
                        bk = str((bi.center.x + res) * 1000 + bi.center.y)
                        bk0 = str((bi.center.x) * 1000 + bi.center.y + res)
                        if self.vertexes.has_key(bk):
                            pflag=True
                            adjb = True
                            bi.edges[bk] = self.vertexes[bk]
                            self.vertexes[bk].edges[str(bi.center.x * 1000 + bi.center.y)] = bi
                    bk = str((bi.center.x) * 1000 + bi.center.y+res)
                    if self.vertexes.has_key(bk):
                        bi.edges[bk] = self.vertexes[bk]
                        self.vertexes[bk].edges[str(bi.center.x * 1000 + bi.center.y)] = bi
                    else:
                        adjb = False
                else:
                    #print bi
                    if adjb == False:
                        firstcon = False
                        bk = str((bi.center.x) * 1000 + bi.center.y - res)
                        if self.vertexes.has_key(bk):
                            adjb = True
                            firstcon = False
                            bi.edges[bk] = self.vertexes[bk]
                            self.vertexes[bk].edges[str(bi.center.x * 1000 + bi.center.y)] = bi
                        bk = str((bi.center.x) * 1000 + bi.center.y + res)
                        bk0 = str((bi.center.x + res) * 1000 + bi.center.y)
                        if self.vertexes.has_key(bk):
                            pflag=True
                            adjb = True
                            bi.edges[bk] = self.vertexes[bk]
                            self.vertexes[bk].edges[str(bi.center.x * 1000 + bi.center.y)] = bi
                    bk = str((bi.center.x + res) * 1000 + bi.center.y )
                    if self.vertexes.has_key(bk):
                        bi.edges[bk] = self.vertexes[bk]
                        self.vertexes[bk].edges[str(bi.center.x * 1000 + bi.center.y)] = bi
                    else:
                        adjb = False
        for k in branches:
            b=branches[k]
            for bi in b:
                if len(bi.edges.keys()) < 2:
                    for i in range(len(instance.agents)):
                        if (not (i==index)) and (instance.agents[i].vertexes.has_key(str(bi.center.x * 1000 + bi.center.y))):
                            try:
                                bi.edges[bi.edges.keys()[0]].edges.pop(str(bi.center.x * 1000 + bi.center.y))
                                self.vertexes.pop(str(bi.center.x * 1000 + bi.center.y))
                                b.remove(bi)
                            except:
                                pass
                if index==1 and (bi.center.x==4.5 or bi.center.x==5.5) and (bi.center.y==4.5 ):
                    bi.edges[bi.edges.keys()[0]].edges.pop(str(bi.center.x * 1000 + bi.center.y))
                    self.vertexes.pop(str(bi.center.x * 1000 + bi.center.y))
                    b.remove(bi)
        self.branches=branches


    def register(self):
        return instance.register(self)

    def directionestm(self, dir):
        return (np.floor(dir/2), dir%2)

    def getremoteleaf(self):
        dis=0
        k=None
        rn = self.unreachedneighbor()
        if len(rn)>0:
            for rk in rn:
                r=rn[rk]
                d = 0
                for a in instance.agents:
                    if a.id == self.id:
                        continue
                    d = d + (a.position-r.center).abs()
                if d >dis:
                    dis=d
                    k=str(r.center.x*1000+r.center.y)
        else:
            c=instance.graph.graphcenter()
            dis = 100000
            for v in self.treeneighbors:
                d = (c-self.treeneighbors[v].center).abs()
                if d < dis:
                    dis=d
                    k=str(self.treeneighbors[v].center.x*1000+self.treeneighbors[v].center.y)
        return self.treeneighbors[k]

    def spantree(self):
         v=self.getremoteleaf()
         ev = self.setposition(v.center)
         return ev

    def gentrj(self,vertex):
        ipos = copy.deepcopy(self.initposition)
        agentpos=self.initposition
        trj = []
        trj.append([ipos.x, ipos.y])
        i=0
        while(abs(ipos.x-agentpos.x)+abs(ipos.y-agentpos.y)>0 or len(trj)<3):
            agentpos,vertex=vertex.step(agentpos)
            trj.append([agentpos.x,agentpos.y])
            i=i+1
            if(i>1000):
                break
        return trj



class swarm:
    def __init__(self):
        self.agents = []

    def copygraph(self,graph):
        self.graph = graph

    def register(self, agent):
        self.agents.append(agent)
        self.agentcount = len(self.agents)
        #print "Now register agent %d"%self.agentcount
        return self.agentcount-1

    def states(self):
        return self.agents

    def agentforce(self, id):
        force=vector2d()
        grc = self.graph.graphcenter()
        cforce = grc-self.agents[id].position
        force=force + cforce/(10.0*cforce.norm())
        for i in range(self.agentcount):
            if i==id:
                continue
            aforce = (self.agents[id].position-self.agents[i].position)
            aforce = vector2d(0 if aforce.x == 0 else 1.0 / aforce.x, 0 if aforce.y == 0 else 1.0 / aforce.y)
            force = force + aforce
        force = force/force.abs()
        print [force.x,force.y]
        print [round(force.x), round(force.y)]
        return vector2d(round(force.x), round(force.y))


    def initagents(self,randmod=True):
        keys = self.graph.vertexes.keys()
        num = len(keys)
        kr= np.random.randint(num,size=len(self.agents))

        i=0
        ipos=[]
        initpos=[]
        ratio=10.0
        initpos.append(vector2d(0.15*ratio,0.05*ratio))
        initpos.append(vector2d(0.75*ratio,0.05*ratio))
        initpos.append(vector2d(0.45*ratio,0.05*ratio))
        initpos.append(vector2d(0.45*ratio,0.85*ratio))
        
        for a in self.agents:
            if randmod:
                a.setposition(self.graph.vertexes[keys[kr[i]]].center)
            else:
                a.setposition(initpos[i]) #fix init position
                print initpos[i].x
            ipos.append([a.initposition.x,a.initposition.y])
            i=i+1
        return np.asarray(ipos)

    def spanonce(self):
        lp = []
        pp = []
        for a in self.agents:
            if instance.graph.isfullycovered():
                break
            ev=a.spantree()
            lp.append(ev)
            pp.append(a.position)
        for a in self.agents:
            for v in a.vertexes:
                if a.vertexes[v].overreached and len(a.vertexes[v].edges)<2:
                    a.vertexes[v].edges[0].edges.pop(v)
                    a.vertexes[v].overreached=False
                    a.vertexes.pop(v)

        return lp, pp

    def sorttree(self):
        pp=[]
        for k in range(len(self.agents)):
            print "Agent:"+str(k)+"/"+str(len(self.agents))
            a=self.agents[k]
            a.sortvertex(index=k)
            ap=[]
            for i in a.branches:
                b = a.branches[i]
                for j in range(len(b)):
                    ap.append(b[j])
            pp.append(ap)
        print "jcoount:"+(str(k))
        return pp




    def gentrj(self):
        trj=[]
        smoothtrj=[]
        for a in self.agents:
            t=np.asarray(a.gentrj(a.initvertex))
            trj.append(t)

        for i in range(len(self.agents)):
            t=trj[i]
            st=[]
            st.append(t[0])
            startvel=0
            for j in range(1,len(t)-1):
                b=bt(startvel=startvel)
                dr=[]
                dr.append(np.asarray(t[j]-st[-1]))
                dr.append(np.asarray(t[j+1]-t[j]))
                #print dr,t[j]
                b.transpattern(t[j],dr)
                ss=b.interpolate()
                #ss=b.evalcurve(100)
                st.extend(ss)
                startvel=b.exitvel
                #print "%d:(%f,%f)"%(i,st[-1][0],st[-1][1])
            smoothtrj.append(st)

        return trj,smoothtrj

instance = swarm()
